#include "key_generator.h"
#include <string>

key_generator::key_generator(){}

key_generator::~key_generator(){}

key_generator::key_generator(QString Dir){
    obj_Dir.setPath(Dir);
    obj_Dir.mkdir("KG");
    obj_Dir.cd("KG");
    //obj_Dir.setPath(Dir.append("/").append("KG"));

    create_key("k");

    // Crea la base de datos que hace la funcion de VotingKeyList(VKL)
    if(!QFile::exists(obj_Dir.absoluteFilePath("VotingKeyList").append(".db"))){
        QString cmd("sqlite3 -line ");
        cmd.append(obj_Dir.absoluteFilePath("VotingKeyList")).append(".db ");
        cmd.append("'create table VotingKeyList(ev BLOB, ez BLOB, dz BLOB);'");
        system(cmd.toAscii().data());
    }

    // Crea la base de datos que hace la funcion de KeyGeneratorBulletinBoard(KGBB)
    if(!QFile::exists(obj_Dir.absoluteFilePath("KGBB").append(".db"))){
        QString cmd("sqlite3 -line ");
        cmd.append(obj_Dir.absoluteFilePath("KGBB")).append(".db ");
        cmd.append("'create table KGBB(Hez BLOB, Hdz BLOB);'");
        system(cmd.toAscii().data());
    }

    if(!QSqlDatabase::contains("KGBBconex"))
        QSqlDatabase::addDatabase("QSQLITE","KGBBconex");
    if(!QSqlDatabase::contains("VKLconex"))
        QSqlDatabase::addDatabase("QSQLITE","VKLconex");
    if(!QSqlDatabase::contains("SVKLconex"))
        QSqlDatabase::addDatabase("QSQLITE","SVKLconex");
}

key_generator::key_generator(QString Dir, bool no_key){
    obj_Dir.setPath(Dir.append("/KG"));
    if(!QSqlDatabase::contains("KGBBconex"))
        QSqlDatabase::addDatabase("QSQLITE","KGBBconex");
    if(!QSqlDatabase::contains("VKLconex"))
        QSqlDatabase::addDatabase("QSQLITE","VKLconex");
    if(!QSqlDatabase::contains("SVKLconex"))
        QSqlDatabase::addDatabase("QSQLITE","SVKLconex");
}

void key_generator::SubVotingKeyList(){
// Crea la conexion
    QSqlDatabase VKL = QSqlDatabase::database("VKLconex",false);
    VKL.setDatabaseName(obj_Dir.absoluteFilePath("VotingKeyList.db"));
    if(!VKL.open()){
        return;
    }

// Crea el modelo
    QSqlTableModel VKLmodel(0,VKL);
    VKLmodel.setTable("VotingKeyList");
    VKLmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    VKLmodel.select();

// Crea la base de datos que hace la funcion de SubVotingKeyList
    if(QFile::exists(obj_Dir.absoluteFilePath("SubVotingKeyList.db"))){
        QFile::remove(obj_Dir.absoluteFilePath("SubVotingKeyList.db"));
    }

    QString cmd("sqlite3 -line ");
    cmd.append(obj_Dir.absoluteFilePath("SubVotingKeyList.db "));
    cmd.append("'create table SubVotingKeyList(ev BLOB, dz BLOB);'");
    system(cmd.toAscii().data());


    QSqlDatabase SVKL = QSqlDatabase::database("SVKLconex",false);
    SVKL.setDatabaseName(obj_Dir.absoluteFilePath("SubVotingKeyList.db"));

    if(!SVKL.open()){
        return;
    }

// Crea el modelo
    QSqlTableModel SVKLmodel(0,SVKL);
    SVKLmodel.setTable("SubVotingKeyList");
    SVKLmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    SVKLmodel.select();

    for(int i = 0;i<VKLmodel.rowCount();i++){
        SVKLmodel.insertRows(0,1);
        SVKLmodel.setData(SVKLmodel.index(0,SVKLmodel.fieldIndex("ev")),VKLmodel.record(i).value("ev"));
        SVKLmodel.setData(SVKLmodel.index(0,SVKLmodel.fieldIndex("dz")),VKLmodel.record(i).value("dz"));
        SVKLmodel.submitAll();
    }
}

bool key_generator::M2(QString &m2,QString &m4){
/*  Desencripta m2
    Verifica Db
    Desencripta el mensaje firmado para obtener la llave publica del votante, y la guarda
    Crear un par de llaves de votacion (ez,dz)
    Guardar las llaves en VotingKeyList (ev,ez,dz)
    Publicar los Hashes H(ev,ez) y H(ev,dz) en KGBB
    Generar M3
        --abrir ez (llave publica)
        --firmar ez
        --agregar ev al mensaje firmado de ez
        --encriptar lo anterior con la llave publica recibida del votante
    Generar M4
        -- abre la llave publica del KG
        -- une el mensaje M3 con la llave publica del KG
        -- Firma la union anterior
        -- Encripta lo anterior con la llave publica del BG*/

    QString Ek,Db, Ekk, eb, ev,Dk;
    QString hexKey,signedKey,m3;
    int Ekk_size;

    // Desencripta m2 y guarda el resultado en Ek
    if(!decryptMsg(m2,Ek)){
        return false;
    }

    // Verifica la firma del BG en Ek y guarda el mensaje en Db; además divide el contenido del último en Ekk y eb

    if(!verifyMsg(Ek,"eb",Db)){
        return false;
    }

    Ekk_size = shift_String(Db).toInt();
    Ekk = shift_String(Db,Ekk_size);
    eb=Db;

    // Verifica que la integridad del mensaje al chequear eb
    if(!Key2Hex("eb",hexKey)){
        return false;
    }

    if(eb!=hexKey){
        return false;
    }

    // Desencripta Ekk y guarda la llave publica del votante(ev)
    if(!decryptMsg(Ekk,ev)){
        return false;
    }

    if(!Hex2Key(ev,obj_Dir.absoluteFilePath("ev"))){
        return false;
    }

    // Crea las llaves del votante
    create_key("z");
    // Guarda ev, ez y dz en VKL.db
    if(!Save_in_VKL()){
        return false;
    }

    // Guarda los hashes H(ev,ez) y H(ev,dz);
    if(!Save_Hashes()){
        return false;
    }

    // Genera M3
    if(!Key2Hex("ez",hexKey)){
        return false;
    }

    signedKey = signMsg(hexKey);

    int signed_size = signedKey.size();
    signedKey.prepend("&&");
    signedKey.prepend(QString::number(signed_size));
    signedKey.append(ev);

    m3 = encryptMsg(signedKey,"ev");

    // Genera M4
    if(!Key2Hex("ek",hexKey)){
        return false;
    }

    int m3_size = m3.size();
    m3.prepend("&&");
    m3.prepend(QString::number(m3_size));
    m3.append(hexKey);

    Dk = signMsg(m3);
    m4 = encryptMsg(Dk,"eb");

    return true;
}

void key_generator::create_key(QString key_name){
    QString namePubKey="e";
    QString namePriKey="d";

    namePubKey.append(key_name);
    namePriKey.append(key_name);

    InvertibleRSAFunction params;
    params.GenerateRandomWithKeySize( rng, 1024);
    RSA::PrivateKey pri( params );

    RSA::PublicKey pub(pri);

    SavePublicKey(namePubKey,pub);
    SavePrivateKey(namePriKey,pri);
}

void key_generator::SavePrivateKey(QString &filename, const PrivateKey &key){
    ByteQueue queue;
    key.Save(queue);

    Save(filename, queue);
}

void key_generator::SavePublicKey(QString &filename, const PublicKey &key){
    ByteQueue queue;
    key.Save(queue);

    Save(filename, queue);
}

void key_generator::Save(QString &filename, const BufferedTransformation &bt){
    filename.prepend("/");
    filename.prepend(obj_Dir.absolutePath());

    FileSink file(filename.toAscii().data());

    bt.CopyTo(file);
    file.MessageEnd();
}

void key_generator::LoadPublicKey(QString filename, PublicKey &key){
    // http://www.cryptopp.com/docs/ref/class_byte_queue.html
    ByteQueue queue;
    Load(filename, queue);
    key.Load(queue);
}

void key_generator::LoadPrivateKey(QString filename, PrivateKey &key){
    // http://www.cryptopp.com/docs/ref/class_byte_queue.html
    ByteQueue queue;

    Load(filename, queue);
    key.Load(queue);
}

void key_generator::Load(QString filename, BufferedTransformation& bt)
{
    // http://www.cryptopp.com/docs/ref/class_file_source.html
    filename.prepend("/");
    filename.prepend(obj_Dir.absolutePath());

    FileSource file(filename.toAscii().data(), true /*pumpAll*/);
    file.TransferTo(bt);
    bt.MessageEnd();
}

QString key_generator::encryptMsg(QString Data,QString keyName){
    QString out;
    string plainMsg,encryptedMsg,buf;

//Convierto de QString a std::string
//En estos métodos es importantísimo porque si se usa QString directamente, corta el mensaje y despues no se puede desencriptar
    plainMsg=Data.toStdString();

    RSA::PublicKey pub;

    LoadPublicKey(keyName,pub);

    RSAES_OAEP_SHA_Encryptor e( pub);

    while(plainMsg.size()>0){
        buf.clear();
        if(plainMsg.size()<=86){
            buf=plainMsg;
            plainMsg.clear();

        }
        else{
            for(int j=0;j<86; j++){
                buf+=plainMsg.at(j);
            }
            plainMsg.erase(0,86);
        }

        StringSource(buf, true,
            new PK_EncryptorFilter( rng, e,
                new StringSink( encryptedMsg )
            ) // PK_EncryptorFilter
         ); // StringSource
    }

    return QString::fromStdString(encryptedMsg);
}

bool key_generator::decryptMsg(QString Data, QString &decryptedMsg){
    QString plainMsg;

    string decrypted_Msg,buf;
    RSA::PrivateKey pri;

    LoadPrivateKey("dk",pri);

    RSAES_OAEP_SHA_Decryptor d( pri );

    //Convierto de QString a std::string
    //En estos métodos es importantísimo porque si se usa QString directamente, corta el mensaje y no se puede desencriptar
    try{
        while(Data.size()>0){
            buf=QString(shift_String(Data,128)).toStdString();

            StringSource(buf, true,
                new PK_DecryptorFilter( rng, d,
                    new StringSink( decrypted_Msg )
                ) // PK_DecryptorFilter
             ); // StringSource
        }
    }
    catch(...){
        return false;
    }

    decryptedMsg = QString::fromStdString(decrypted_Msg);
    return true;
}

QString key_generator::Hash(QString keyname1, QString keyname2){
// Abre el archivo key1 y carga su contenido en un QByteArray BA_key1
    QFile file1(obj_Dir.absoluteFilePath(keyname1));
    file1.open(QIODevice::ReadOnly);
    QByteArray BA_key1 = file1.readAll();
    file1.close();

// Abre el archivo key2 y carga su contenido en un QByteArray BA_key2
    QFile file2(obj_Dir.absoluteFilePath(keyname2));
    file2.open(QIODevice::ReadOnly);
    QByteArray BA_key2 = file2.readAll();
    file2.close();

// Genera la funcion hash de las llaves key1 y key2
    SHA sha;
    string output;

    sha.Update((byte const*)BA_key1.data(), BA_key1.size());
    sha.Update((byte const*)BA_key2.data(), BA_key2.size());

    SecByteBlock sbbDigest(sha.DigestSize());

    sha.Final(sbbDigest.begin());
    HexEncoder enc(new StringSink(output));
    enc.Put(sbbDigest.begin(), sbbDigest.size());
    enc.MessageEnd();

    return QString::fromStdString(output);
}

QString key_generator::signMsg(QString msg){
    string plainMsg,signedMsg;

//Convierto de QString a std::string
//En estos métodos es importantísimo porque si se usa QString directamente, corta el mensaje y despues no se puede desencriptar
    plainMsg=msg.toStdString();

// Carga la llave privada del objeto
    RSA::PrivateKey pri;
    LoadPrivateKey("dk",pri);

    RSASS<PSSR, SHA>::Signer signer(pri);

    StringSource(plainMsg, true,
        new SignerFilter(rng, signer,
            new StringSink(signedMsg),
            true // putMessage for recovery
        ) // SignerFilter
     ); // StringSource*/

    return QString::fromStdString(signedMsg);
}

bool key_generator::verifyMsg(QString msg,QString keyToUse,QString &verifiedMsg){
    string signedMsg, verified_Msg;

    signedMsg=msg.toStdString();

    RSA::PublicKey pub;
    LoadPublicKey(keyToUse,pub);

//    CryptoPP::RSASSA_PKCS1v15_SHA_Verifier verifier(pub);
    RSASS<PSSR, SHA>::Verifier verifier(pub);

    try{
        StringSource(signedMsg, true,
            new SignatureVerificationFilter(
                verifier,
                new StringSink(verified_Msg),
                SignatureVerificationFilter::THROW_EXCEPTION | SignatureVerificationFilter::PUT_MESSAGE
           ) // SignatureVerificationFilter
        ); // StringSource
    } catch(Exception e){
        return false;
    }
    verifiedMsg=QString::fromStdString(verified_Msg);

return true;
}

QString key_generator::shift_String(QString &msg){
    QString buffer;

    int shift = msg.indexOf("&&");

    for(int j=0;j<shift; j++){
        buffer.append(msg.at(j));
    }

    msg.remove(0,shift);
    msg.remove(0,2); // borra &&
    return buffer;
}

QString key_generator::shift_String(QString &msg,int max){
    QString fer("errol");

    fer.clear();
    for(int j=0;j<max; j++){
        fer.append(msg.at(j));
    }

    msg.remove(0,max);
    return fer;
}

bool key_generator::Save_in_VKL(){
/*  Crear la conexion con la base de datos
    Crear el modelo que manipule la base de datos
    Abrir los archivos que contienen las llaves VoterKey(ev) y VotingKeyPair(ez y dz)
    Copiar los archivos en la base de datos
  */
// Crea la conexion
    QSqlDatabase VKL = QSqlDatabase::database("VKLconex",false);
    VKL.setDatabaseName(obj_Dir.absoluteFilePath("VotingKeyList.db"));

    if(!VKL.open()){
        return false;
    }

// Crea el modelo
    QSqlTableModel VKLmodel(0,VKL);
    VKLmodel.setTable("VotingKeyList");
    VKLmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    VKLmodel.select();

// Abre las llaves
    QString ev,ez,dz;

    if(!Key2Hex("ev",ev)){
        return false;
    }

    if(!Key2Hex("ez",ez)){
        return false;
    }

    if(!Key2Hex("dz",dz)){
        return false;
    }

// Copia las llaves en la base de datos
    int row = 0;
    VKLmodel.insertRows(row, 1);
    VKLmodel.setData(VKLmodel.index(row, VKLmodel.fieldIndex("ev")),ev);
    VKLmodel.setData(VKLmodel.index(row, VKLmodel.fieldIndex("ez")),ez);
    VKLmodel.setData(VKLmodel.index(row, VKLmodel.fieldIndex("dz")),dz);
    VKLmodel.submitAll();
    return true;
}

bool key_generator::Key2Hex(QString keyname, QString &keyHex){
    QFile file(obj_Dir.absoluteFilePath(keyname));
    if(!file.open(QIODevice::ReadOnly)){
        return false;}

    QByteArray BA = file.readAll();

    string encodedKey;

    HexEncoder hex(new StringSink(encodedKey));
    hex.Put((const byte*)BA.data(),BA.size());
    hex.MessageEnd();

    keyHex=QString::fromStdString(encodedKey);

    return true;
}

bool key_generator::Save_Hashes(){
/*  Crear la conexion con la base de datos
    Crear el modelo que manipule la base de datos
    Crear los siguientes hashes H(ev,ez) y H(ev,dz)
    Copiar los hashes en la base de datos
  */

        // Crea la conexion

        QSqlDatabase KGBB = QSqlDatabase::database("KGBBconex",false);
        KGBB.setDatabaseName(obj_Dir.absoluteFilePath("KGBB.db"));


        if(!KGBB.open()){
            return false;
        }

        // Crea el modelo
        QSqlTableModel KGBBmodel(0,KGBB);
        KGBBmodel.setTable("KGBB");
        KGBBmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
        KGBBmodel.select();

        // Crea los hashes
        QString Hez = Hash("ev","ez");
        QString Hdz = Hash("ev","dz");

        // Copia las llaves en la base de datos
        int row = 0;
        KGBBmodel.insertRows(row, 1);
        KGBBmodel.setData(KGBBmodel.index(row, KGBBmodel.fieldIndex("Hez")),Hez);
        KGBBmodel.setData(KGBBmodel.index(row, KGBBmodel.fieldIndex("Hdz")),Hdz);
        KGBBmodel.submitAll();
    return true;
}

bool key_generator::Hex2Key(QString keydata, QString keyname){
    QFile file(obj_Dir.absoluteFilePath(keyname));
    if(!file.open(QIODevice::WriteOnly)){
        return false;}

//    QByteArray BA = file.readAll();

    string decodedKey;

    HexDecoder deco(new StringSink(decodedKey));
    deco.Put((byte*)keydata.toStdString().data(),keydata.size());
    deco.MessageEnd();

    file.write(decodedKey.data(),decodedKey.size());

    file.close();
    return true;
}

bool key_generator::Search_in_KGBB(QString Hash,QString column){
    QSqlDatabase KGBB = QSqlDatabase::database("KGBBconex",false);
    KGBB.setDatabaseName(obj_Dir.absoluteFilePath("KGBB.db"));

    if(!KGBB.open()){
        return false;
    }

    // Crea el modelo
    QSqlTableModel KGBBmodel(0,KGBB);
    KGBBmodel.setTable("KGBB");
    KGBBmodel.setEditStrategy(QSqlTableModel::OnManualSubmit);
    KGBBmodel.select();

    // verifica si Hash se encuentra en la base de datos
    for (int i = 0; i <= KGBBmodel.rowCount(); ++i) {
        if (Hash==KGBBmodel.record(i).value(column).toString()){
            return true;
        }
    }
    return false;
}






